#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @python: 3.6

import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader


def test_img(net_g, datatest, args):
    net_g.eval()
    # testing
    test_loss = 0
    correct = 0
    data_loader = DataLoader(datatest, batch_size=args.bs)
    l = len(data_loader)
    # data_loader 是一个 PyTorch DataLoader 对象，可以被用于遍历数据集。
    # enumerate(data_loader) 返回一个 (idx, (data, target)) 元组，其中 idx 是当前数据的索引，(data, target) 包含了当前的输入数据和标签。
    for idx, (data, target) in enumerate(data_loader):
        # 这里判断是否需要使用 GPU 进行计算，如果 args.gpu 不等于 -1，则将输入数据和标签转移到 GPU 上进行计算
        if args.gpu != -1:
            data, target = data.cuda(), target.cuda()
        # net_g 是神经网络结构，这里将输入数据 data 作为输入，得到输出结果 log_probs
        log_probs = net_g(data)
        # sum up batch loss
        # 使用交叉熵损失函数 F.cross_entropy 计算当前批次的损失值，并加到 test_loss 变量上，reduction='sum' 表示对每个样本的损失值进行求和
        test_loss += F.cross_entropy(log_probs, target, reduction='sum').item()
        # get the index of the max log-probability
        # 根据输出结果 log_probs 得到预测标签 y_pred
        # log_probs.data 表示输出结果 log_probs 中的数据。max(1, keepdim=True) 表示沿着第二个维度（即列）取最大值，保持维度不变。
        # [1] 表示返回索引值
        y_pred = log_probs.data.max(1, keepdim=True)[1]
        # 计算当前批次预测正确的样本数，并加到 correct 变量上，.long() 表示将 True/False 转换为整数类型。
        # .cpu().sum() 表示将计算结果转移到 CPU 上并求和
        correct += y_pred.eq(target.data.view_as(y_pred)).long().cpu().sum()

    test_loss /= len(data_loader.dataset)
    accuracy = 100.00 * correct / len(data_loader.dataset)
    if args.verbose:
        print('\nTest set: Average loss: {:.4f} \nAccuracy: {}/{} ({:.2f}%)\n'.format(
            test_loss, correct, len(data_loader.dataset), accuracy))
    return accuracy, test_loss

